from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from sklearn.feature_extraction.text import TfidfVectorizer
import pandas as pd
import joblib
import os

# 加载数据
data = pd.read_excel('emotion_data.xlsx')
data = data.dropna(subset=['评价分词'])

# 使用 TF-IDF 进行文本向量化
vectorizer = TfidfVectorizer(max_features=1000)
X = vectorizer.fit_transform(data['评价分词'])
y = data['情感标签']

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 定义模型
models = {
    'SVM': SVC(),
    'DecisionTree': DecisionTreeClassifier(),
    'LogisticRegression': LogisticRegression(),
    'NaiveBayes': MultinomialNB(),
    'KNeighbors': KNeighborsClassifier()
}

# 确保输出目录存在
output_dir = 'output'

# 训练和评估模型
results = {}
for name, model in models.items():
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    results[name] = {
        'Accuracy': accuracy_score(y_test, y_pred),
        'F1-score': f1_score(y_test, y_pred, average='weighted'),
        'Precision': precision_score(y_test, y_pred, average='weighted'),
        'Recall': recall_score(y_test, y_pred, average='weighted')
    }
    # 导出模型
    model_path = os.path.join(output_dir, f"{name}_model.joblib")
    joblib.dump(model, model_path)

# 打印结果
for name, metrics in results.items():
    print(f"Model: {name}")
    for metric, value in metrics.items():
        print(f"{metric}: {value:.4f}")
    print()
